% % ----------------3.12------------------------------------
% THIS IS THE MAIN FILE FOR PRODUCING THE FIGURES AND NUMERICAL RESULTS IN
% THE PAPER
% THE ONLY VERSION OF THE MODEL NOT COMPUTED HERE IS THE POVERTY TRAP
% VERSION WITH INVESTMENT CONSTRAINTS
%-----------------------------------------------------
 clear;

%pause on;

%-------------------------------------------------------
%Model Parameterization
%-------------------------------------------------------

%Indicator for Cases
% THE POLITICAL ECONOMY SPECIFICATION, SECTION 6F IN THE PAPER 
Polecon         = 'N';
% THE ALTERNATIVE CALIBRATION IN THE SECOND LINE OF TABLE 1 
Alternative_Cal = 'N';
% THE SPECIFICATION WITH FIXED INVESTMENT, SECTION 6C IN THE PAPER
Irreversible    = 'N';

%------------------
%Basic parameters
%------------------

eta     = 1.0;          % parameter of the exponential for cost distribution (normalization?)

%Externally set values
UStau   = 0.181;        % US progressivity rate 2000-2006: 0.181 FROM PSID, 0.20 from CBO with adjustment; 
sigma   = 2.0;          % curvature of disutility of work (from framework) 
rstar   = 0.04;         % Interest rate
g       = 0.189;        % fraction of output devoted to (valued) expenditures;

%JK:these estimates change because of changed tau
%see Mathematica NB get_vars.wls for the calculation
%old values are commented out

%From Framework estimates (average 2000-2006)
veps    = 0.177; %0.164;        % cross-sectional variance of insurable wage shocks 
theta   = 3.245; %3.124;        % elasticity of substitution in production & 1/theta^2 is cons. ineq. at age 0
vphi    = 0.028; %0.036;        % cross-sectional variance of preference heterogeneity    


%Cross-sectional residual uninsurable inequality
valpha  = 0.093; %0.098;        % valpha = valpha0+delta/(1-delta)*vomega;
valpha0  = 0;           % Initial residual uninsurable inequality = zero       
delta   = 0.971;        % survival rate

psi = 0.65;                % New parameter defining elasticity of skill investment

nul = 0.5;                 % Low value for inequality aversion parameter
nuh = 2;                   % High value for inequality aversion parameter

%-----------------------------
%Original draft values for basic parameters
%-----------------------------
%eta     = 1.0;          % parameter of the exponential for cost distribution (normalization?)
%UStau   = 0.151;        % US progressivity rate 2000-2006: 0.151 FROM PSID, 0.20 from CBO with adjustment; 
%sigma   = 2.165;        % curvature of disutility of work (from framework) 
%rstar   = 0.04;         % Interest rate
%g       = 0.189;        % fraction of output devoted to (valued) expenditures;
%veps    = 0.166;        % cross-sectional variance of insurable wage shocks 
%theta   = 3.144;        % elasticity of substitution in production & 1/theta^2 is cons. ineq. at age 0
%vphi    = 0.035;        % cross-sectional variance of preference heterogeneity    
%valpha  = 0.097;        % valpha = valpha0+delta/(1-delta)*vomega;
%valpha0  = 0;           % Initial residual uninsurable inequality = zero       
%delta   = 0.971;        % survival rate
%psi = 1;                % New parameter defining elasticity of skill investment

%------------------
%Derivative parameters
%------------------

taum1 = UStau;          % Define the tau inherited from the past  


chi   = g/(1-g);        % weight on utility from G (used to be 0.25)
                        % chi = 0.057;   value for chi that rationalizes tau = 0.151

vomega = (valpha-valpha0)*(1-delta)/delta;

rho   = rstar +(1-UStau)*(2-UStau)*vomega/2;    % Rate of time preference
beta  = 1/(1+rho);                              % Discount factor
gamma = beta;                                   % Weight the planner uses to discount future generations


%valpha=0;
%valpha0=0;
%vomega=0;
%vphi=0;
%veps=0;

% Following using to compute G in the US given g and UStau

%UStau=0.084;


sigmahat_UStau=(sigma+UStau)/(1-UStau);

logYUStau = (theta/(theta-1))*log(theta/(theta-1))+((psi/((theta-1)*(1+psi)))+(1/(1+sigma)))*log(1-UStau)-(psi/((theta-1)*(1+psi)))*log(theta)-(1/((theta-1)*(1+psi)))*log(eta)+...
    ((1/sigmahat_UStau^2)*(UStau+sigmahat_UStau+sigmahat_UStau*UStau)*(veps/2));

Y_UStau = exp(logYUStau);
           
agg_G   = g*Y_UStau;        % Needed to compute average output in the baseline optimum



%Alternative calibration with theta = 2 based on Pareto tail from PSID
if Alternative_Cal == 'Y'

    theta   = 2.000;        % elasticity of substitution in production & 1/theta^2 is cons. ineq. at age 0
    vphi    = 0.021;        % cross-sectional variance of preference heterogeneity    
    veps    = 0.139;        % cross-sectional variance of insurable wage shocks 
    
    vunins0 = (1/theta)^2;  % cross-sectional variance of uninsurable initial conditions 
    valpha   = 0.0;
    valpha0  = 0; 
    vomega  = 0.000;        % variance of uninsurable life-cycle shocks 

    rho   = rstar +(1-UStau)*(2-UStau)*vomega/2;
    beta = 1/(1+rho);
    gamma = beta;           % Define the weight the planner uses to discount future generations

    psi = 1;                % New parameter defining elasticity of skill investment
    
    logYUStau = (theta/(theta-1))*log(theta/(theta-1))+((psi/((theta-1)*(1+psi)))+(1/(1+sigma)))*log(1-UStau)-(psi/((theta-1)*(1+psi)))*log(theta)-(1/((theta-1)*(1+psi)))*log(eta)+...
    ((1/sigmahat_UStau^2)*(UStau+sigmahat_UStau+sigmahat_UStau*UStau)*(veps/2));

    Y_UStau = exp(logYUStau);

    agg_G   = g*Y_UStau;     % need to recompute average output in the baseline optimum  
 
end
    


%Decomposition of SWF
for i=1:1999;
    
    tau(i)=-1+0.001*i;

    sigmahat(i)=(sigma+tau(i))/(1-tau(i));
    
    %Aggregate Output
 
    logY = (theta/(theta-1))*log(theta/(theta-1))+((psi/((theta-1)*(1+psi)))+(1/(1+sigma)))*log(1-tau(i))-(psi/((theta-1)*(1+psi)))*log(theta)-(1/((theta-1)*(1+psi)))*log(eta)+...
    ((1/sigmahat(i)^2)*(tau(i)+sigmahat(i)+sigmahat(i)*tau(i))*(veps/2));

    Y(i) = exp(logY);
    
    %M term in allocations
    Mscript(i)=( (1-tau(i))*(1-tau(i)*(1+sigmahat(i)))/sigmahat(i) )*(veps/2);
    
    %Base-price for skills
    pi0(i) = (1/(theta-1))*log(theta/(theta-1))+(psi/((theta-1)*(1+psi)))*log((1-tau(i))/theta)-(1/((theta-1)*(1+psi)))*log(eta);

    %Integral of y_{i}^{1-tau} across agents (w/o lambda term)
    Yd_1(i)=(1-tau(i))^((1-tau(i))/(1+sigma))*exp(-(1/sigmahat(i))*Mscript(i))*exp(((1-tau(i))*(1+sigmahat(i))/sigmahat(i))*( ((1-tau(i))*(1+sigmahat(i))/sigmahat(i)) - 1 )*(veps/2)); 
    Yd_2(i)= exp((1-tau(i))*pi0(i))*exp(-tau(i)*(1-tau(i))*vphi/2);
    Yd_3(i)=((1-delta)*exp(-tau(i)*(1-tau(i))*(valpha0/2)) / (1-delta*exp((-tau(i)*(1-tau(i))/2)*vomega)))*(theta/(-1+tau(i)+theta));
    
    Yd(i)=Yd_1(i)*Yd_2(i)*Yd_3(i);
    
    %Lambda from aggregate resource which is used in the case where G is
    %non-valued and aggregate G is fixed (case 3)
    lambda3(i) = (Y(i)-agg_G)/Yd(i);
    
    %SWF with chi=0 and an exogenous amount G wasted (case 3)
    W_TOT_XG3(i) =     log(lambda3(i))+((1-tau(i))/(1+sigma))*log(1-tau(i))+Mscript(i)-(1-tau(i))/(1+sigma)+...                   % this line has terms in log(c) (except unins. wage) and disutility from hours  
      (1-tau(i))*pi0(i)+(1-tau(i))/theta-(1-tau(i))*(valpha0/2)-(1-tau(i))*(delta/(1-delta))*(vomega/2)-(1-tau(i))*(vphi/2)-...   % this line has E[(1-tau)*(log(p)+alpha-phi)]
      (psi/(1+psi))*(1-tau(i))/theta+(((1-beta)*delta)/(1-beta*delta))*(psi/(1+psi))*(1/theta)*(1-taum1); ;                      % this line has education cost  

    %SWF with chi=0 and an exogenous amount G wasted for  right-wing planner (case 3)
    
    % SWF of the representative agent with chi=0 (and g=0)
    W_RA_NOG(i)    = ( 1/(1+sigma) )*( log(1-tau(i)) - (1-tau(i)) );
    
    % SWF of the representative agent with valued G
    % Note that log(1+g)+chi log g = chi*log(chi) - (1+chi)*log(1+chi)
    W_RA(i)   = W_RA_NOG(i) + chi*log(chi) - (1+chi)*log(1+chi) + chi*( log(1-tau(i))/(1+sigma) );
    
    %SWF piece associated to productivity gain from skill investment
    %Second line is for the case chi=0 (and g=0)
    W_EDU1(i)      = (1+chi)*( (psi/(1+psi))*(1/(theta-1))*log(1-tau(i)) + (1/((theta-1)*(1+psi)))*log( (1/(eta*theta^psi))*(theta/(theta-1))^(theta*(1+psi)) ) );
    W_EDU1_NOG(i) = ( (psi/(1+psi))*(1/(theta-1))*log(1-tau(i)) + (1/((theta-1)*(1+psi)))*log( (1/(eta*theta^psi))*(theta/(theta-1))^(theta*(1+psi)) ) );
    
    %SWF piece associated to between skill-group consumption inequality
    %(second line for right-wing planner)
    W_EDU2(i)      =  (1-tau(i))/theta  + log(1-((1-tau(i))/theta));

    W_EDU2_NUL(i)      =  log(1-((1-tau(i))/theta)) - (nul/(nul-1))*log(((1-nul)/nul)*(1/theta)*(1-tau(i))+1);
    W_EDU2_NUH(i)      =  log(1-((1-tau(i))/theta)) - (nuh/(nuh-1))*log(((1-nuh)/nuh)*(1/theta)*(1-tau(i))+1);
    W_EDU2_NUINF(i)    =  0;

    %SWF piece associated to average education cost (second line for right-wing planner)
    W_EDU3(i)      =     -(psi/(1+psi))*(1/theta)*(1-tau(i)) + (((1-beta)*delta)/(1-beta*delta))*(psi/(1+psi))*(1/theta)*(1-taum1);
        
    %SWF piece associated to preference heterogeneity (second line for right-wing planner)
    W_PREF(i)      = -(1-tau(i))^2*vphi/2;

    W_PREF_NUL(i)      = -(1-tau(i))^2*(1/nul)*vphi/2;
    W_PREF_NUH(i)      = -(1-tau(i))^2*(1/nuh)*vphi/2;
    W_PREF_NUINF(i)    = 0;
    
    %SWF piece associated with within skill-group consumption inequality
    W_OMEGA(i)      = log( ( 1-delta*exp( -tau(i)*(1-tau(i))*vomega/2 ) )/(1-delta) ) - (1-tau(i))*delta/(1-delta)*vomega/2;
%   W_OMEGA_APPX(i) = log( ( 1-delta*exp( -tau(i)*(1-tau(i))*vomega/2 ) )/(1-delta) ) - tau(i)*(1-tau(i))*delta/(1-delta)*vomega/2;
    W_ALPHA_APPROX(i) = -(1-tau(i))^2*valpha/2;
    
 
    W_OMEGA_NUCOMMON(i) = -(1-tau(i))^2*((beta*delta)/(1-beta*delta))*(vomega/2); 

    W_OMEGA_NUL(i) = -(1-tau(i))^2*(1/nul)*((1-beta)/(1-beta*delta))*(delta/(1-delta))*(vomega/2); 
    W_OMEGA_NUH(i) = -(1-tau(i))^2*(1/nuh)*((1-beta)/(1-beta*delta))*(delta/(1-delta))*(vomega/2); 
    W_OMEGA_NUINF(i) = 0; 
    
    %SWF piece associated to initial variance of uninsuranle shocks
    W_ALPHA0(i) = -(1-tau(i))^2*valpha0/2;

    %SWF piece associated to insurable shocks (second line for chi=0 (and g=0))
    W_EPS(i)  = (1+chi)*(1-tau(i))/(sigma+tau(i))*veps - (1+chi)*sigma*( (1-tau(i))/(sigma+tau(i)) )^2*veps/2 ;
    W_EPS_NOG(i)  = (1-tau(i))/(sigma+tau(i))*veps - sigma*( (1-tau(i))/(sigma+tau(i)) )^2*veps/2 ;
    
    %SWC piece associated to insurable shock for progressive consumption tax
    W_CONSTAX_EPS(i)  = ((1+chi)/sigma)*veps - (1+chi)*sigma*(1/sigma^2)*(veps/2) ;

    %Additional terms for SWF defined as ex-ante expected welfare of agent
    %with median and average values of (kappa, alpha0, phi) before you draw your epsilon
    
    W_ADJ(i) =           (1-tau(i))*(vphi/2) + (1-tau(i))*(delta/(1-delta))*(vomega/2) - (1-tau(i))/theta + (psi/(1+psi))*((1-tau(i))/theta - (1-taum1)/theta) - (beta*delta/(1-beta*delta))*(1-tau(i))*(vomega/2);  % Here we undo the expectations of the idiosyncratic terms in welfare
    
    W_TOT_MEAN_ADD(i) = -(1-tau(i))*(vphi/2) - (1-tau(i))*(delta/(1-delta))*(vomega/2) + (1-tau(i))/theta - (psi/(1+psi))*((1-tau(i))/theta - (1-taum1)/theta);
    W_TOT_MED_ADD(i) =  -(1-tau(i))*(vphi/2) - (1-tau(i))*(delta/(1-delta))*(vomega/2) + (1-tau(i))*log(2)/theta - (psi/(1+psi))*((1-tau(i))*log(2)/theta - (1-taum1)*log(2)/theta);
    
       
end;

%-----------------------------------------
%Aggregation of components of SWF
%and calculation of optimal taus
%----------------------------------------

%Piece together edu terms (second line is right-wing planner)
W_EDU   = W_EDU1 + W_EDU2 + W_EDU3;

W_EDU_NUL = W_EDU1 + W_EDU2_NUL + W_EDU3;
W_EDU_NUH = W_EDU1 + W_EDU2_NUH + W_EDU3;
W_EDU_NUINF = W_EDU1 + W_EDU2_NUINF + W_EDU3;

%Piece together edu terms in case chi=0 (and g=0) (second line is right-wing planner)
W_EDU_NOG   = W_EDU1_NOG + W_EDU2 + W_EDU3;

%Piece together initial dispersion and lifetime shocks
W_ALPHA = W_ALPHA0+W_OMEGA;

%Piece together all components (second line is approximation of lifetime cumulation of Vomega)
W_TOT = W_RA + W_EDU + W_PREF + W_ALPHA + W_EPS;
W_TOT_APPROX = W_RA+W_EDU+W_PREF+W_ALPHA_APPROX+W_EPS;

W_TOT_NUL = W_RA + W_EDU_NUL + W_PREF_NUL + W_OMEGA_NUCOMMON + W_OMEGA_NUL + W_EPS;
W_TOT_NUH = W_RA + W_EDU_NUH + W_PREF_NUH + W_OMEGA_NUCOMMON + W_OMEGA_NUH + W_EPS;
W_TOT_NUINF = W_RA + W_EDU_NUINF + W_PREF_NUINF + W_OMEGA_NUCOMMON + W_OMEGA_NUINF + W_EPS;

%SWF in the case where skills are exogenous
W_TOT_NOSKILL2 = W_TOT - W_EDU + W_EDU2; 

%SWF for progressive consumption tax
W_CONSTAX_TOT = W_RA+W_EDU+W_PREF+W_ALPHA+W_CONSTAX_EPS;

%Ex-ante SWF for agent with median and mean (kappa, phi,alpha)
W_TOT_MED = W_TOT + W_ADJ + W_TOT_MED_ADD;
W_TOT_MEAN = W_TOT + W_ADJ + W_TOT_MEAN_ADD;

%Various SWF for case when chi=0 (and g=0)
W_TOT_NOG       = W_RA_NOG + W_EDU_NOG + W_PREF + W_ALPHA + W_EPS_NOG;
%SWF for case where chi=0 and an exogenous fraction of output g is wasted 
W_TOT_XG2     = W_TOT + log(1-g) - chi*W_EPS/(1+chi) - chi*W_EDU1/(1+chi) + (W_RA_NOG - W_RA);


%-----------------------------
%Compute implied tax system
%-----------------------------

% Representative Agent
disp('RA taustar');
[c,i]=max((W_RA)');
optau_RA = tau(i)

% Piece associated to skills investment
disp('Skill Investment taustar');
[c,i]=max((W_EDU)');
optau_EDU = tau(i)

% Optimal tau
disp('Exact taustar');
[c,i]=max(W_TOT');
optau_EXACT = tau(i)
disp('Welfare gain from tax reform with HSV tau value');
100*(exp(max(W_TOT')-W_TOT(1181))-1)

disp('Welfare gain from tax reform with actual tau');
100*(exp(max(W_TOT')-W_TOT(1129))-1)


